use std::io;
use std::io::{Error, Read, Write};
use std::net::{TcpListener, TcpStream};
use std::thread;
use std::time::Instant;

fn handle_client(mut stream: TcpStream) -> Result<(), Error> {
    let mut buf: [u8; 1024] = [0; 1024];

    loop {
        stream.read_exact(&mut buf[0..1])?;

        if buf[0] == 0 {
            stream.shutdown(std::net::Shutdown::Both)?;
            break;
        } else if buf[0] == 1 {
            stream.write(&[1])?;
        } else if buf[0] == 2 {
            stream.write(&buf)?;
        } else if 20 >= buf[0] && buf[0] > 10 {
            for _ in 0..(buf[0] - 10) {
                for _ in 0..1024 {
                    stream.write(&buf)?;
                }
            }
        } else if 30 >= buf[0] && buf[0] > 20 {
            stream.write(&[3])?;
            for _ in 0..(buf[0] - 20) {
                for _ in 0..1024 {
                    stream.read_exact(&mut buf)?;
                }
            }
            stream.write(&[4])?;
        }
    }

    Ok(())
}

fn server() -> std::io::Result<()> {
    let mut ipaddr = String::new();

    println!("Please input the bind port:");
    io::stdin().read_line(&mut ipaddr)?;

    let listener = TcpListener::bind("0.0.0.0:".to_string()+ipaddr.trim())?;
    
    let mut thread_vec: Vec<thread::JoinHandle<()>> = Vec::new();

    for stream in listener.incoming() {
        let stream = stream.expect("failed!");
        let handle = thread::spawn(move || {
            handle_client(stream).unwrap_or_else(|error| eprint!("{:?}", error));
        });

        thread_vec.push(handle);
    }

    for handle in thread_vec {
        handle.join().unwrap();
    }

    Ok(())
}

fn average(array: &[f32]) -> f32 {
    let mut count: f32 = 0.0;
    for i in 1..11 {
        count += array[i - 1] / (i as f32);
    }
    count /= array.len() as f32;

    count
}

fn send_data(
    len: &usize,
    stream: &mut TcpStream,
    buf: &mut [u8],
    ttl: &mut f32,
    send_time: &mut f32,
) -> std::io::Result<()> {
    // compute ttl
    let ttl_timer = Instant::now();
    stream.write(&[1])?; // send 1
    stream.read_exact(&mut buf[0..1])?;
    *ttl = 0.0005 * (ttl_timer.elapsed().as_micros() as f32);

    // send data
    stream.write(&[(*len as u8) + 20])?;
    loop {
        stream.read_exact(&mut buf[0..1])?;
        if buf[0] == 3 {
            break;
        }
    }
    let recv_start = Instant::now();
    for _ in 0..*len {
        for _ in 0..1024 {
            stream.write(buf)?;
        }
    }
    loop {
        stream.read_exact(&mut buf[0..1])?;
        if buf[0] == 4 {
            break;
        }
    }
    *send_time = 0.001 * (recv_start.elapsed().as_micros() as f32) - *ttl;

    Ok(())
}

fn recv_data(
    len: &usize,
    stream: &mut TcpStream,
    buf: &mut [u8],
    ttl: &mut f32,
    recv_time: &mut f32,
) -> std::io::Result<()> {
    // compute ttl
    let ttl_timer = Instant::now();
    stream.write(&[1])?; // send 1
    stream.read_exact(&mut buf[0..1])?;
    *ttl = 0.0005 * (ttl_timer.elapsed().as_micros() as f32);

    // send data
    stream.write(&[(*len as u8) + 10])?;
    let recv_start = Instant::now();

    for _ in 0..*len {
        for _ in 0..1024 {
            stream.read_exact(buf)?;
        }
    }
    *recv_time = 0.001 * (recv_start.elapsed().as_micros() as f32) - *ttl;

    Ok(())
}

fn client() -> std::io::Result<()> {
    println!("Please input the server's addr(ip:port):");

    let mut ipaddr = String::new();
    io::stdin().read_line(&mut ipaddr)?;

    let mut stream = TcpStream::connect(&ipaddr.trim())?;

    let mut buf: [u8; 1024] = [0; 1024];

    const ROUTINE: usize = 10;

    let mut recv_time = [0.0; ROUTINE];
    let mut send_time = [0.0; ROUTINE];

    for i in 1..(ROUTINE + 1) {
        let mut ttl = 0.0;

        recv_data(&i, &mut stream, &mut buf, &mut ttl, &mut recv_time[i - 1])?;

        println!(
            "recv {:02}MB {:.03}ms\t{:.03}MBps\tttl {:.03}ms",
            i,
            recv_time[i - 1],
            (i as f32) * 1e3 / recv_time[i - 1],
            ttl
        );
    }

    let aver_recv = average(&recv_time);
    println!("=============================================");
    println!(
        "recv {:.03} ms per MB w.r.t. {:.03} MBps",
        aver_recv,
        1e3 / aver_recv
    );
    println!("=============================================");

    for i in 1..(ROUTINE + 1) {
        let mut ttl = 0.0;

        send_data(&i, &mut stream, &mut buf, &mut ttl, &mut send_time[i - 1])?;

        println!(
            "send {:02}MB {:.03}ms\t{:.03}MBps\tttl {:.03}ms",
            i,
            send_time[i - 1],
            (i as f32) * 1e3 / send_time[i - 1],
            ttl
        );
    }

    let aver_send = average(&send_time);
    println!("=============================================");
    println!(
        "send {:.03} ms per MB w.r.t. {:.03} MBps",
        aver_send,
        1e3 / aver_send
    );
    println!("=============================================");

    stream.write(&[0])?;

    Ok(())
}

fn main() -> std::io::Result<()> {
    println!("Work as\n 1) server \n 2) client\npress the number to continue:");

    let mut input = String::new();

    io::stdin().read_line(&mut input)?;

    let status: u8 = input.trim().parse().expect("Please input num");

    if status == 1 {
        // call server
        server()?;
    } else if status == 2 {
        // call client
        client()?;
    } else {
        // call error
    }

    Ok(())
}
